import torch
import torchvision.transforms.functional as TF
from PIL import Image
from pathlib import Path
import matplotlib.pyplot as plt
from __init__ import UNet

# 加载 U-Net 模型（确保已定义并训练好）
model = UNet(in_channels=3, out_channels=1)
model.load_state_dict(torch.load(r"D:\Data\Project\code_learn\annotated_deep_learning_paper_implementations-master\annotated_deep_learning_paper_implementations-master\logs\unet\81bbffdc0ed611f08905001a7dda7113\checkpoints\10176\model.pth", map_location=torch.device('cpu')))
model.eval()

# 你的测试图片文件夹
image_path = Path(r"D:\Data\Project\code_learn\annotated_deep_learning_paper_implementations-master\annotated_deep_learning_paper_implementations-master\labml_nn\unet\test_images")
output_path = Path(r"D:\Data\Project\code_learn\annotated_deep_learning_paper_implementations-master\annotated_deep_learning_paper_implementations-master\labml_nn\unet\output_masks")
output_path.mkdir(exist_ok=True)

# 处理测试图片
def preprocess_image(image: Image.Image):
    """
    预处理输入图片：转换为 Tensor
    """
    image_tensor = TF.to_tensor(image).unsqueeze(0)  # 添加 batch 维度
    return image_tensor


# 进行推理并保存结果
def predict_and_save(image_file: Path):
    """
    读取图片 -> 预测掩码 -> 调整尺寸 -> 保存掩码
    """
    # 如果是 PNG 格式，转为 JPG 格式
    if image_file.suffix.lower() == '.png':
        image = Image.open(image_file).convert("RGB")
    else:
        image = Image.open(image_file).convert("RGB")

    original_size = image.size[::-1]  # 记录原始尺寸 (H, W)

    # 预处理
    image_tensor = preprocess_image(image)

    # 预测
    with torch.no_grad():
        mask = model(image_tensor)

    # 变换回图片格式
    mask = mask.squeeze(0).squeeze(0)  # 去掉 batch 维度和通道维度
    mask = mask.unsqueeze(0)  # 变成 (1, H, W) 以便 resize
    mask = TF.resize(mask, original_size)  # 调整为原图尺寸
    mask = mask.squeeze(0)  # 变回 (H, W)
    mask = (mask > 0.5).float()  # 阈值化处理
    mask_image = TF.to_pil_image(mask)

    # 保存掩码
    output_file = output_path / f"mask_{image_file.stem}.png"
    mask_image.save(output_file)
    print(f"Saved: {output_file}")

    # 显示结果
    fig, ax = plt.subplots(1, 2, figsize=(8, 4))
    ax[0].imshow(image)
    ax[0].set_title("Original Image")
    ax[1].imshow(mask_image, cmap="gray")
    ax[1].set_title("Predicted Mask")
    plt.show()


# 处理所有测试图片
for img_file in image_path.glob("*"):
    predict_and_save(img_file)